#!/usr/bin/env python

# Imports:
# -------
# WX and standard python modules
import wx, os, time, threading

# Project
from labeler_panel import *
from annotation_panel import *
from topic_dialog import *

# ROS
from annotations.msg import Annotation, AnnotationArray


# Constants & configuration:
M_FILE_QUIT_ID = 1
M_FILE_OPEN_ID = 2
M_FILE_CLOSE_ID = 3
M_FILE_SAVE_ID = 4
M_FILE_SAVE_AS_ID = 5
M_HELP_ABOUT_ID = 6

# -----------------------
class MainFrame(wx.Frame):
# -----------------------
	# ----------------------------------------
	def __init__(self, bagLabeler, name, size):
	# ----------- ----------------------------
		# Initialze the frame (no resizing)
		wx.Frame.__init__(self, None, wx.ID_ANY, name, size = size, \
						  style = wx.MINIMIZE_BOX | wx.MAXIMIZE_BOX | wx.SYSTEM_MENU \
							 	| wx.CAPTION | wx.CLOSE_BOX | wx.CLIP_CHILDREN)
		
		# Save the labeler
		self.bagLabeler = bagLabeler

		# Setup the menu bar
		self.menuBar = wx.MenuBar()
		self.SetMenuBar(self.menuBar)

		# File menu:
		# ---------
		mFile = wx.Menu()
		self.menuBar.Append(mFile, '&File')

		# Open bag
		mFileOpen = wx.MenuItem(mFile, M_FILE_OPEN_ID, '&Open Bag\tCtrl+O')
		self.Bind(wx.EVT_MENU, self.onOpen, id=M_FILE_OPEN_ID)
		mFile.AppendItem(mFileOpen)

		# Close bag
		mFileClose = wx.MenuItem(mFile, M_FILE_CLOSE_ID, '&Close Bag\tCtrl+C')
		self.Bind(wx.EVT_MENU, self.onBagClose, id=M_FILE_CLOSE_ID)
		mFile.AppendItem(mFileClose)

		# Save annotations
		self.mFileSave = wx.MenuItem(mFile, M_FILE_SAVE_ID, '&Save Annotations\tCtrl+S')
		self.Bind(wx.EVT_MENU, self.onSave, id=M_FILE_SAVE_ID)
		mFile.AppendItem(self.mFileSave)
		self.mFileSave.Enable(False)

		# 'Save as' annotations
		self.mFileSaveAs = wx.MenuItem(mFile, M_FILE_SAVE_AS_ID, 'Save Annotations As')
		self.Bind(wx.EVT_MENU, self.onSaveAs, id=M_FILE_SAVE_AS_ID)
		mFile.AppendItem(self.mFileSaveAs)
		self.mFileSaveAs.Enable(False)

		# Quit
		mFileQuit = wx.MenuItem(mFile, M_FILE_QUIT_ID, '&Quit\tCtrl+Q')
		self.Bind(wx.EVT_MENU, self.onQuit, id=M_FILE_QUIT_ID)
		mFile.AppendItem(mFileQuit)

		# Help menu:
		# ---------
		mHelp = wx.Menu()
		mHelpAbout = wx.MenuItem(mHelp, M_HELP_ABOUT_ID, '&About')
		self.Bind(wx.EVT_MENU, self.onAbout, id=M_HELP_ABOUT_ID)
		mHelp.AppendItem(mHelpAbout)
		self.menuBar.Append(mHelp, '&Help')

		# Create the sizer
		self.sizer = wx.BoxSizer(wx.HORIZONTAL)
		self.SetSizer(self.sizer)

		# Create and add the annotation panel
		# -----------------------------------
		self.sizer.AddSpacer(10)
		self.annotationPanel = AnnotationPanel(self)
		self.sizer.Add(self.annotationPanel, 1, wx.EXPAND)
		self.Bind(wx.EVT_BUTTON, self.onNewAnnotation, self.annotationPanel.newButton)
		self.Bind(wx.EVT_BUTTON, self.onDeleteAnnotation, self.annotationPanel.deleteButton)
		self.Bind(wx.EVT_LISTBOX, self.onSelectAnnotation, self.annotationPanel.listBox)

		# Create and add the labeler panel
		# --------------------------------
		self.sizer.AddSpacer(10)
		self.labelerPanel = LabelerPanel(self)	
		self.labelerPanel.setRubberBandCallback(self.onRubberBand)
		self.sizer.Add(self.labelerPanel, 3, wx.EXPAND)
		self.Bind(wx.EVT_BUTTON, self.onLeft, self.labelerPanel.leftButton)
		self.Bind(wx.EVT_BUTTON, self.onRight, self.labelerPanel.rightButton)
		self.Bind(wx.EVT_CHAR_HOOK, self.leftRightDispatch)
		self.Bind(wx.EVT_SLIDER, self.onSlider)
		self.Bind(wx.EVT_TEXT, self.onNameText, self.labelerPanel.nameField)
		self.Bind(wx.EVT_COMBOBOX, self.onOrientationSelect, self.labelerPanel.orientationCombo)
		self.labelerPanel.SetFocus()
		self.sizer.AddSpacer(10)
		
		# Wire up the close event
		self.Bind(wx.EVT_CLOSE, self.onClose)

		# Put everything together
		self.Layout()

	# ---------------------
	def onOpen(self, event):
	# ---------------------
		# Open a dialog box for choosing (opening a file)
		openDlg = wx.FileDialog(self, "Choose a ROS Bag", "~", "", "*.bag", wx.FD_OPEN)
		if openDlg.ShowModal() == wx.ID_OK:
			fileName = openDlg.GetFilename()
			dirName = openDlg.GetDirectory()
			fullPath = os.path.join(dirName, fileName)
			openDlg.Destroy()

			# Show the "loading..." message
			wx.CallAfter(self.labelerPanel.setImage, self.labelerPanel.loadingBag)
				
			# Load the bag
			def loadBag():
				try:
					self.bagLabeler.loadBag(fullPath, callback=self.onBagLoad)
				except Exception as e:
					wx.MessageBox('Error loading bag file: %s' % str(e), 'Error')
					self.onQuit()
			wx.CallAfter(loadBag)
		else:
			openDlg.Destroy()

	# -------------------------
	def onBagClose(self, event):
	# -------------------------
		print 'TODO: implement bag closing'

	# ------------------
	def save(self, path):
	# ------------------
		self.bagLabeler.annotationOutPath = path
		self.bagLabeler.saveAnnotations(path)

	# ---------------------
	def onSave(self, event):
	# ---------------------
		outPath = self.bagLabeler.annotationOutPath
		if self.bagLabeler.annotationOutPath != None:
			self.save(outPath)
		else:
			self.onSaveAs(event)

	# -----------------------
	def onSaveAs(self, event):
	# -----------------------
		saveDlg = wx.FileDialog(self, "Save as", "~", "", "*.bag", wx.FD_SAVE)
		if saveDlg.ShowModal() == wx.ID_OK:
			fileName = saveDlg.GetFilename()
			fileName = fileName if fileName.endswith('.bag') else fileName + '.bag'
			dirName = saveDlg.GetDirectory()
			fullPath = os.path.join(dirName, fileName)
			saveDlg.Destroy()

			# Save!
			self.save(fullPath)
		else:
			saveDlg.Destroy()

	# ------------------------------------------------------------------
	def onBagLoad(self, imageTopics, annotationTopics, cameraInfoTopics):
	# ------------------------------------------------------------------
		def matchInfoTopic(iTopic):
			baseList = iTopic.rpartition('/')
			base = baseList[0]
	
			cit = base + '/' + 'camera_info' # TODO: This shouldn't have to be a hard constant
			if cit in cameraInfoTopics:
				return cit
			else: return None

		# Open the dialog box 
		topicDialog = TopicDialog(self, wx.ID_ANY, imageTopics, annotationTopics)
		topicDialog.ShowModal() 

		# Determing if the user pressed the 'ok' button (or if they aborted by hitting the x)
		if topicDialog.okPressed:
			topicDialog.Destroy()

			# Process the dialog box 
			annotationTopic = topicDialog.annotationTopic
			imageTopic = topicDialog.imageTopic
			rectify = topicDialog.rectify
			shuffle = topicDialog.shuffle

			# New annotation topic
			# --------------------
			if annotationTopic == '[New annotation topic]':
				# Get the new topic's name
				inputDialog = wx.TextEntryDialog(self, 'Annotation topic name:', 'Name annotation topic', 'ground_truth')			
				while True:
					if inputDialog.ShowModal() == wx.ID_OK:
						annotationTopic = inputDialog.GetValue()
						break
				wx.CallAfter(self.bagLabeler.setImageTopic, imageTopic, rectify, shuffle,  matchInfoTopic(imageTopic))
				wx.CallAfter(self.bagLabeler.newAnnotationTopic, annotationTopic)
				
			# Annotation topic from external bag
			# ----------------------------------
			elif annotationTopic == '[Load external annotation bag]':
				# Get the bag path
				loadDlg = wx.FileDialog(self, "Load annotation bag", "~", "", "*.bag", wx.FD_OPEN)
				annotationBagPath = None
				while True:
					if loadDlg.ShowModal() == wx.ID_OK:
						fileName = loadDlg.GetFilename()
						dirName = loadDlg.GetDirectory()
						annotationBagPath = os.path.join(dirName, fileName)
						loadDlg.Destroy()
						break
				wx.CallAfter(self.bagLabeler.setImageTopic, imageTopic, rectify, shuffle, matchInfoTopic(imageTopic))
				def setAnnotation():
					result = self.bagLabeler.loadExternalAnnotations(annotationBagPath, shuffle)
					if not result:
						print 'NO ANNOTATION TOPICS' # TODO: Do better than this
				wx.CallAfter(setAnnotation)

			# Annotation topic from this bag
			# ------------------------------
			else:
				log.error('Illegal annotation source: %s' % str(annotationTopic))
				return False

			def finishSetup():
				# Enable the vcr controls on the labeler panel
				self.labelerPanel.enableVCR()
				
				# Enable the annotation panel
				self.annotationPanel.Enable()

				# Enable saving
				self.mFileSave.Enable(True)
				self.mFileSaveAs.Enable(True)
			
				# Get the first image and display it
				topic, img, t, aa = self.bagLabeler.curr()
				self.update(topic, img, t, aa)

				# Indicate that everything is OK
			wx.CallAfter(finishSetup)
			return True
	
		else:
			# This will get tripped if the user exits out of the dialog, so we want to abort loading
			self.labelerPanel.setImage(self.labelerPanel.home)
			self.labelerPanel.Refresh()
			return False


	# ----------------------------
	def onQuit(self, event = None):
	# ----------------------------
		""" Called when the 'Quit' button is pressed in the menu """
		self.Close()

	# ----------------------
	def onAbout(self, event):
	# ----------------------
		""" Launches an 'about' dialog """

		info = wx.AboutDialogInfo()
		info.SetIcon(wx.Icon('resources/bag.png', wx.BITMAP_TYPE_PNG))
		info.SetName('Bag Labeler')
		info.SetVersion('0.1')
		info.SetDescription('A simple tool for labeling ROS bags.\n Written by Walker Burgin, Feb. 2011')
		info.SetWebSite('http://www.ros.org')
		info.AddDeveloper('Walker Burgin - <walkerburgin@gmail.com>')
		wx.AboutBox(info)
		
	# ----------------------
	def onClose(self, event):
	# ----------------------
		""" Called when the application is closed """ 

		# Make sure the user really wants to quit
		dial = wx.MessageDialog(self, 'You have unsaved annotations, are you sure to quit?', 'Question', 
			wx.YES_NO | wx.NO_DEFAULT | wx.ICON_QUESTION)
		if not self.bagLabeler.unsavedWork or dial.ShowModal() == wx.ID_YES:

			# Shutdown the bag labeler
			self.bagLabeler.shutdown()

			# Destroy the GUI window
			self.Destroy()

	# -----------------------
	def onSlider(self, event):
	# -----------------------
		# Get the image
		val = self.labelerPanel.slider.GetValue()
		topic, img, t, aa = self.bagLabeler.setIndex(val)
		self.update(topic, img, t, aa)

	# ---------------------
	def onLeft(self, event):
	# ---------------------
		if self.bagLabeler.isBagLoaded():
			# Get & display the next image
			topic, img, t, aa = self.bagLabeler.prev()
			self.update(topic, img, t, aa)		
	
	# ----------------------
	def onRight(self, event):
	# ----------------------
		if self.bagLabeler.isBagLoaded():
			# Get & display the next image
			topic, img, t, aa = self.bagLabeler.next()
			self.update(topic, img, t, aa)			

	# -----------------------------------------------------
	def update(self, topic, image, time, aa, selection = 0):
	# -----------------------------------------------------
		# Set the image
		self.labelerPanel.setImage(image)

		# Display the progress
		i, tot = self.bagLabeler.getIndex()
		self.labelerPanel.displayProgress(i,tot)

		# Display the annotations
		if aa != None:
			name, orientation = self.annotationPanel.showAnnotations(aa, selection)
			self.labelerPanel.nameField.SetValue(name)
			if orientation != '':
				self.labelerPanel.orientationCombo.SetStringSelection(orientation)

			# Enable/disable the new/delete buttons & the labeler controls
			if len(aa.annotations) > 0:
				self.annotationPanel.deleteButton.Enable()
				self.labelerPanel.enableLabeler()
		
				# Draw the annotations
				foregroundRects = []
				backgroundRects = []
				for i, a in enumerate(aa.annotations):
					# Put together the rectangle
					rect = (aa.annotations[i].x,
							aa.annotations[i].y,
							aa.annotations[i].width, 
							aa.annotations[i].height)

					if i == selection:
						foregroundRects.append(rect)
					else:
						backgroundRects.append(rect)

				# Draw the rectangles
				self.labelerPanel.rubberBand.drawEvt(foregroundRects, backgroundRects)

			else:
				self.annotationPanel.deleteButton.Disable()
				self.labelerPanel.disableLabeler()

		# All done? refresh!
		self.labelerPanel.rubberBand.Refresh()
		
	# --------------------------------
	def leftRightDispatch(self, event):
	# --------------------------------
		kc = event.GetKeyCode()
		if kc == wx.WXK_RIGHT:
			self.onRight(event)
		elif kc == wx.WXK_LEFT: 
			self.onLeft(event) 
		else:
			event.Skip()
	
	# --------------------------
	def onRubberBand(self, rect):
	# --------------------------
		# Set the annotation rectangle
		topic, img, t, aa = self.bagLabeler.curr()
		i = self.annotationPanel.getSelectionIndex()
		if i != -1:
			aa.annotations[i].x         	= rect[0]
			aa.annotations[i].y      		= rect[1]
			aa.annotations[i].width  		= rect[2]
			aa.annotations[i].height 		= rect[3]

		# Indicate that we've got unsaved work
		self.bagLabeler.unsavedWork = True
			
	# ------------------------------
	def onNewAnnotation(self, event):
	# ------------------------------
		# Create a new annotation message 
		topic, img, t, aa = self.bagLabeler.curr()

		annotation = Annotation()
		annotation.x = 0
		annotation.y = 0
		annotation.width = 0
		annotation.height = 0
		annotation.name = 'object_%i' % len(aa.annotations)
		annotation.name = annotation.name.encode('ascii', 'ignore') # Required for serialization
		annotation.header = aa.header

		# Add the annotation to the AnnotationArray
		aa.annotations.append(annotation)

		# Update the display
		self.update(topic, img, t, aa, len(aa.annotations) - 1)

	# ---------------------------------
	def onDeleteAnnotation(self, event):
	# ---------------------------------
		# Find the annotation that's selected
		i = self.annotationPanel.getSelectionIndex()
		if i != -1:
			# Remove the annotation message from the annotation array
			topic, img, t, aa = self.bagLabeler.curr()
			del aa.annotations[i]

			# Update the display
			self.update(topic, img, t, aa)

	# -------------------------
	def onNameText(self, event):
	# -------------------------
		val = self.labelerPanel.nameField.GetValue()
		i = self.annotationPanel.getSelectionIndex()
		if i != -1:
			topic, img, t, aa = self.bagLabeler.curr()
			oldName = aa.annotations[i]
			if val != oldName:
				# Update the annotation
				aa.annotations[i].name = val.encode('ascii', 'ignore') 	# Required for serialization
				self.annotationPanel.showAnnotations(aa, i)

	# ---------------------------------
	def onSelectAnnotation(self, event):
	# ---------------------------------
		i = self.annotationPanel.getSelectionIndex()
		if i != -1:
			topic, img, t, aa = self.bagLabeler.curr()

			name = aa.annotations[i].name
			self.labelerPanel.nameField.SetValue(name)
			
			orientation = aa.annotations[i].orientation
			if orientation == '': 
				val = self.labelerPanel.orientationCombo.GetValue()
				aa.annotations[i].orientation = val.encode('ascii', 'ignore') # Required for serialization
			self.labelerPanel.orientationCombo.SetStringSelection(aa.annotations[i].orientation)
			self.update(topic, img, t, aa, i)

	# ----------------------------------
	def onOrientationSelect(self, event):
	# ----------------------------------
		val = self.labelerPanel.orientationCombo.GetValue()
		i = self.annotationPanel.getSelectionIndex()
		topic, img, t, aa = self.bagLabeler.curr()
		aa.annotations[i].orientation = val.encode('ascii', 'ignore') # Required for serialization
